package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"
)

var (
	beforeRank = 10                                  // 爬取榜单前多少位的作品
	dirLoc     = "D:/Pixiv日榜/"                       // 保存目录
	u, err     = url.Parse("http://127.0.0.1:41091") // 代理（这里用的是自己的代理

	aimUrl      = `https://www.pixiv.net/ranking.php?mode=daily&content=illust`
	chanImgUrls chan string // 图片下载链接管道
	chanTask    chan string // uid
	waitGroup   sync.WaitGroup
	rePid       = `<a href="/artworks/\d+"class="work  _work  "target="_blank">`
	reImg       = `https://i.pximg.net/img-original/img[^\"]+?(\.((jpg)|(png)|(jpeg)|(gif)|(bmp)))`

	// 使用代理
	client = http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyURL(u),
		},
	}
)

func main() {
	// 1. 初始化管道
	chanImgUrls = make(chan string, 10) // 图片下载链接管道

	rankPid := GetPid(aimUrl)

	chanTask = make(chan string, beforeRank)

	// 2.爬虫协程 通过pid获取图片链接
	for i := 0; i < len(rankPid); i++ {
		waitGroup.Add(1)
		go GetImgUrls("https://www.pixiv.net/artworks/" + rankPid[i])
	}
	waitGroup.Add(1)
	go checkOk()

	// 3.下载协程
	for i := 0; i < 5; i++ {
		waitGroup.Add(1)
		go DownloadImg()
	}
	waitGroup.Wait()
	fmt.Println("爬取结束！")
}

func checkOk() {
	var count int
	for {
		url := <-chanTask
		fmt.Printf("%s完成了爬取任务\n", url)
		count++
		if count == beforeRank {
			close(chanImgUrls)
			break
		}
	}
	waitGroup.Done()
}

// GetPageStr 获取排行榜页面内容
func GetPageStr(url string) (pageStr string) {
	resp, err := client.Get(url)
	handleError(err, "GetPageStr client.Get")
	defer resp.Body.Close()
	pageBytes, err := ioutil.ReadAll(resp.Body)
	//fmt.Printf("pageBytes:%v",pageBytes)
	handleError(err, "GetPageStr ioutil.ReadAll")
	pageStr = string(pageBytes)
	return
}

// GetPid 获取排行榜整页的pid
func GetPid(url string) (urls []string) {
	pageStr := GetPageStr(url)
	re := regexp.MustCompile(rePid)
	results := re.FindAllStringSubmatch(pageStr, -1)
	rePid := regexp.MustCompile(`\d+`)
	fmt.Printf("共找到%d条结果\n", len(results))
	for _, result := range results {
		url := rePid.FindString(result[0])
		urls = append(urls, url)
		if len(urls) == beforeRank {
			break
		}
	}
	return
}

// GetImgUrls 获取当前pid所对应的图片链接
func GetImgUrls(url string) {
	pageStr := GetPageStr(url)
	reImg := regexp.MustCompile(reImg)
	results := reImg.FindAllStringSubmatch(pageStr, -1) // 图片链接
	chanImgUrls <- results[0][0]
	chanTask <- results[0][0]
	waitGroup.Done()
}

// DownloadImg 下载图片
func DownloadImg() {
	for url := range chanImgUrls {
		urlSp := strings.Split(url, "/")
		fileName := fmt.Sprintf("%s_%s_%s_%s_%s_%s_pid_%s", urlSp[5], urlSp[6], urlSp[7], urlSp[8], urlSp[9], urlSp[10], urlSp[11])
		ok := DownloadFile(url, fileName)
		if ok {
			fmt.Printf("%s下载成功！\n", fileName)
		} else {
			fmt.Printf("%s下载失败！\n", fileName)
		}
	}
	waitGroup.Done()
}

func handleError(err error, why string) {
	if err != nil {
		fmt.Println(err, why)
	}
}

func DownloadFile(url, fileName string) (ok bool) {
	// 由于pixiv的图片访问来源需要是pixiv域名下url，因此需要加一个referer
	req, err := http.NewRequest("GET", url, nil)
	handleError(err, "DownloadFile http.NewRequest")
	req.Header.Add("referer", "https://www.pixiv.net/")

	resp, err := client.Do(req)
	handleError(err, "DownloadFile http.Get")
	defer resp.Body.Close()
	bytes, err := ioutil.ReadAll(resp.Body)
	handleError(err, "DownloadFile ioutil.ReadAll")

	year, month, day := time.Now().Date()
	dirName := fmt.Sprintf("%v年%v月%v日", year, int(month), day)
	mdirerr := os.MkdirAll(dirLoc+dirName, os.ModePerm)
	handleError(mdirerr, "新建文件夹")
	fileName = dirLoc + dirName + "/" + fileName
	//fileName = dirLoc + fileName

	// 写出数据
	err = ioutil.WriteFile(fileName, bytes, 0666)
	if err != nil {
		handleError(err, "文件写出错误，请检查路径是否出现问题")
		return false
	} else {
		return true
	}
}
